#!/usr/bin/env python
# -*-coding: utf-8-*-
"""
@Author:      chen ming
@Date:        2025/5/10 9:46
@Description: 
"""

import numpy as np
from scipy.integrate import odeint


class ROVDynamics:
    def __init__(self, mass=100, volume=0.1):
        # 平移惯性矩阵（质量项）
        self.mass = mass
        self.I_trans = np.eye(3) * mass  # 3x3

        # 旋转惯性矩阵（对角阵）
        self.I_rot = np.diag([10, 10, 5])  # 3x3

        # 完整6x6惯性矩阵（分块对角）
        self.I = np.block([[self.I_trans, np.zeros((3, 3))],
                           [np.zeros((3, 3)), self.I_rot]])

        # 流体动力学参数
        self.Ma = np.diag([50, 50, 50, 10, 10, 10])  # 附加质量矩阵
        self.D_lin = np.diag([20, 20, 20, 5, 5, 5])  # 线性阻尼系数
        self.D_nonlin = 0.1  # 非线性阻尼系数

        # 推进器配置矩阵 (示例为4推进器系统)
        self.B = np.array([
            [1, 1, 0, 0],
            [0, 0, 1, 1],
            [0, 0, 0, 0],
            [0, 0, 0.5, -0.5],
            [0.5, -0.5, 0, 0],
            [0, 0, 0, 0]
        ])

    def coriolis_matrix(self, v):
        # 提取角速度分量
        omega = v[3:]  # 角速度向量 [p, q, r]

        # 构造反对称矩阵（仅处理旋转分量）
        S = np.array([[0, -omega[2], omega[1]],
                      [omega[2], 0, -omega[0]],
                      [-omega[1], omega[0], 0]])

        # 科里奥利矩阵分块构造
        C_rot = -S @ self.I_rot  # 核心修正点：使用旋转惯量矩阵
        return np.block([[np.zeros((3, 3)), np.zeros((3, 3))],
                         [np.zeros((3, 3)), C_rot]])

    def damping_force(self, v):
        # 阻尼力计算（包含非线性项）
        return self.D_lin @ v + self.D_nonlin * np.abs(v) * v


def dynamics_model(state, t, rov, u):
    η = state[:3]   # 位置 (x,y,z)
    θ = state[3:6]  # 欧拉角 (roll,pitch,yaw)
    v = state[6:]   # 线速度/角速度

    # 转换矩阵 J(η)
    R = np.array([[np.cos(θ[2]), -np.sin(θ[2]), 0],
                 [np.sin(θ[2]), np.cos(θ[2]), 0],
                 [0, 0, 1]])
    J = np.block([[R, np.zeros((3,3))],
                 [np.zeros((3,3)), np.eye(3)]])

    # 广义力计算
    τ = rov.B @ u  # 推进器输入转换为广义力

    # 计算加速度
    M_total = rov.I + rov.Ma  # 可正常执行

    C_total = rov.coriolis_matrix(v)
    D_total = rov.damping_force(v)

    v_dot = np.linalg.inv(M_total) @ (τ - C_total @ v - D_total)

    # 状态导数
    η_dot = J @ v
    return np.concatenate((η_dot, v_dot))


def simulate_rov():
    # 验证维度一致性
    rov = ROVDynamics()
    v_sample = np.array([1, 0, 0, 0.5, 0.2, 0.1])  # 线速度+角速度
    C_matrix = rov.coriolis_matrix(v_sample)

    print("C_matrix shape:", C_matrix.shape)  # 应输出(6,6)
    print(C_matrix)
    print("S @ I_rot shape:", (rov.coriolis_matrix(v_sample)[3:, 3:]).shape)  # 应输出(3,3)
    print(rov.coriolis_matrix(v_sample)[3:, 3:])

    t = np.linspace(0, 10, 10000)  # 仿真时间10秒
    u = np.array([10, 10, 10, 10])  # 推进器输入

    # 初始状态 [x,y,z, roll,pitch,yaw, u,v,w,p,q,r]
    state0 = np.zeros(12)

    # 数值积分求解
    states = odeint(dynamics_model, state0, t, args=(rov, u))

    # 可视化位置变化
    import matplotlib.pyplot as plt
    plt.plot(t, states[:, 0], label='X Position')
    plt.plot(t, states[:, 1], label='Y Position')
    plt.plot(t, states[:, 2], label='Depth')
    plt.plot(states[:, 0], states[:, 1], label='X Y')
    plt.legend()
    plt.show()


if __name__ == "__main__":
    simulate_rov()

